#include "cuda_util.h"

__global__ void YUV420P_to_NV12_kernel(
    const uint8_t *__restrict__ u_plane, const uint8_t *__restrict__ v_plane, uint8_t *__restrict__ uv_plane, int uv_width, int uv_height, int u_pitch,
    int v_pitch, int uv_pitch)
{
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    if (x < uv_width && y < uv_height)
    {
        int uv_index = y * uv_pitch + x * 2;
        uv_plane[uv_index + 0] = u_plane[y * u_pitch + x];
        uv_plane[uv_index + 1] = v_plane[y * v_pitch + x];
    }
}

__device__ inline uint8_t clamp(int v)
{
    return v < 0 ? 0 : (v > 255 ? 255 : v);
}

__device__ void BGR_to_YUV(uint8_t b, uint8_t g, uint8_t r, uint8_t &y, uint8_t &u, uint8_t &v)
{
    // BT.709
    y = clamp((47 * r + 157 * g + 16 * b) >> 8);
    u = clamp(((-26 * r - 87 * g + 112 * b) >> 8) + 128);
    v = clamp(((112 * r - 102 * g - 10 * b) >> 8) + 128);
}

extern "C" __global__ void BGR_to_NV12_Kernel(
    const uint8_t *__restrict__ bgr, int bgr_pitch, uint8_t *__restrict__ y_plane, int y_pitch, uint8_t *__restrict__ uv_plane, int uv_pitch, int width,
    int height)
{
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    if (x >= width || y >= height)
        return;

    int bgr_offset = y * bgr_pitch + x * 3;
    uint8_t B = bgr[bgr_offset + 0];
    uint8_t G = bgr[bgr_offset + 1];
    uint8_t R = bgr[bgr_offset + 2];

    uint8_t Y_val, U_val, V_val;
    BGR_to_YUV(B, G, R, Y_val, U_val, V_val);

    y_plane[y * y_pitch + x] = Y_val;

    if ((y % 2 == 0) && (x % 2 == 0))
    {
        int uv_x = x;
        int uv_y = y / 2;

        int sum_u = 0, sum_v = 0;
        for (int dy = 0; dy < 2; dy++)
        {
            for (int dx = 0; dx < 2; dx++)
            {
                int px = x + dx;
                int py = y + dy;
                if (px >= width)
                    px = width - 1;
                if (py >= height)
                    py = height - 1;
                int offset = py * bgr_pitch + px * 3;
                uint8_t Bb = bgr[offset + 0];
                uint8_t Gb = bgr[offset + 1];
                uint8_t Rb = bgr[offset + 2];
                uint8_t y_tmp, u_tmp, v_tmp;
                BGR_to_YUV(Bb, Gb, Rb, y_tmp, u_tmp, v_tmp);
                sum_u += u_tmp;
                sum_v += v_tmp;
            }
        }
        uint8_t avg_u = sum_u / 4;
        uint8_t avg_v = sum_v / 4;

        int uv_offset = uv_y * uv_pitch + uv_x * 2;
        uv_plane[uv_offset + 0] = avg_u;
        uv_plane[uv_offset + 1] = avg_v;
    }
}

void YUV420P_to_NV12(
    const uint8_t *u_plane, const uint8_t *v_plane, uint8_t *uv_plane, int width, int height, int u_pitch, int v_pitch, int uv_pitch, cudaStream_t stream)
{
    dim3 block(16, 16);
    dim3 grid((width + 15) / 16, (height + 15) / 16);

    YUV420P_to_NV12_kernel<<<grid, block, 0, stream>>>(u_plane, v_plane, uv_plane, width, height, u_pitch, v_pitch, uv_pitch);
}

void BGR_to_NV12(const uint8_t *src, uint8_t *y_plane, uint8_t *uv_plane, int src_step, int y_step, int uv_step, int width, int height, cudaStream_t stream)
{
    dim3 block(16, 16);
    dim3 grid((width + 15) / 16, (height + 15) / 16);

    BGR_to_NV12_Kernel<<<grid, block, 0, stream>>>(src, src_step, y_plane, y_step, uv_plane, uv_step, width, height);
}
